from celery import shared_task
from django.core.cache import cache
from .models import ZhihuComment
from .serializers import ZhihuCommentSerializer

@shared_task
def build_comment_tree(task_id, page, page_size):
    try:
        # 定义缓存键
        comment_tree_cache_key = f"zhihu_comment_tree_{task_id}_{page}_{page_size}"
        task_status_key = f"zhihu_comment_tree_status_{task_id}_{page}_{page_size}"

        # 获取评论数据
        comments = ZhihuComment.objects.filter(content_id__task_id=task_id)

        # 构建评论字典
        comment_dict = {comment.comment_id: comment for comment in comments}

        # 构建评论树
        comment_tree = []
        children_map = {}
        for comment in comment_dict.values():
            parent_id = comment.parent_comment_id
            if parent_id != '0':
                children_map.setdefault(parent_id, []).append(comment)
            else:
                comment_tree.append(comment)

        # 分页处理
        start_index = (page - 1) * page_size
        end_index = start_index + page_size
        paginated_comment_tree = comment_tree[start_index:end_index]

        # 序列化评论树
        serialized_comment_tree = ZhihuCommentSerializer(
            paginated_comment_tree, 
            many=True, 
            context={'children': children_map}
        ).data

        # 缓存评论树数据
        cache.set(comment_tree_cache_key, serialized_comment_tree, timeout=86400)
        # 更新任务状态为完成
        cache.set(task_status_key, 'completed', timeout=86400)

        return True
    except Exception as e:
        # 如果发生错误，更新任务状态
        cache.set(task_status_key, f'failed: {str(e)}', timeout=86400)
        return False